#include <asm/mmpage_type.h>
#include <asm/mmpage.h>
#include <mm.h>
#include <asm/mmpage_pro.h>
#include <asm/base.h>
#include <sysregs.h>
#include <string.h>
#include <stdint.h>
#include <stdbool.h>
#define NO_BLOCK_MAPPINGS BIT(0)
#define NO_CONT_MAPPINGS BIT(1)

extern char idmap_pg_dir[];

extern char _text_boot[], _etext_boot[];
extern char _text[], _etext[];


/*debug*/

// 简化的PTE打印函数 - 合并连续映射
static void print_pte_level_simple(pmd_t *pmdp, unsigned long addr, unsigned long end)
{
    pte_t *ptep = pte_offset_phys(pmdp, addr);
    unsigned long range_start = addr;
    unsigned long last_pa = 0;
    unsigned long last_prot = 0;
    int count = 0;
    bool first_mapping = true;
    
    do {
        pte_t pte = *ptep;
        if (!pte_none(pte)) {
            unsigned long pte_pa = pte_val(pte) & PAGE_MASK;
            unsigned long pte_prot = pte_val(pte) & ~PAGE_MASK;
            
            if (count == 0) {
                // 开始第一个映射
                range_start = addr;
                last_pa = pte_pa;
                last_prot = pte_prot;
                count = 1;
            } else if (pte_pa == last_pa + PAGE_SIZE && pte_prot == last_prot) {
                // 连续映射，继续计数
                count++;
                last_pa = pte_pa;
            } else {
                // 映射不连续，打印之前的范围
                if (count == 1) {
                    printk("    PTE: VA=0x%lx -> PA=0x%lx, prot=0x%lx\n",
                           range_start, last_pa, last_prot);
                } else {
                    printk("    PTE: VA=0x%lx-0x%lx -> PA=0x%lx-0x%lx, prot=0x%lx [%d pages]\n",
                           range_start, range_start + (count-1)*PAGE_SIZE,
                           last_pa - (count-1)*PAGE_SIZE, last_pa, last_prot, count);
                }
                // 开始新的范围
                range_start = addr;
                last_pa = pte_pa;
                last_prot = pte_prot;
                count = 1;
            }
        } else {
            // 遇到未映射的PTE，结束当前范围
            if (count > 0) {
                if (count == 1) {
                    printk("    PTE: VA=0x%lx -> PA=0x%lx, prot=0x%lx\n",
                           range_start, last_pa, last_prot);
                } else {
                    printk("    PTE: VA=0x%lx-0x%lx -> PA=0x%lx-0x%lx, prot=0x%lx [%d pages]\n",
                           range_start, range_start + (count-1)*PAGE_SIZE,
                           last_pa - (count-1)*PAGE_SIZE, last_pa, last_prot, count);
                }
                count = 0;
            }
        }
    } while (ptep++, addr += PAGE_SIZE, addr != end);
    
    // 打印最后一个范围
    if (count > 0) {
        if (count == 1) {
            printk("    PTE: VA=0x%lx -> PA=0x%lx, prot=0x%lx\n",
                   range_start, last_pa, last_prot);
        } else {
            printk("    PTE: VA=0x%lx-0x%lx -> PA=0x%lx-0x%lx, prot=0x%lx [%d pages]\n",
                   range_start, range_start + (count-1)*PAGE_SIZE,
                   last_pa - (count-1)*PAGE_SIZE, last_pa, last_prot, count);
        }
    }
}
// 打印PTE级页表内容
static void print_pte_level(pmd_t *pmdp, unsigned long addr, unsigned long end)
{
    pte_t *ptep = pte_offset_phys(pmdp, addr);
    unsigned long i = 0;
    unsigned long current_addr = addr;
    
    printk("PTE Level [0x%lx - 0x%lx]:\n", addr, end);
    do {
        pte_t pte = *ptep;
        if (!pte_none(pte)) {
            unsigned long pte_pa = (pte_val(pte) & PAGE_MASK);
            unsigned long pte_prot = pte_val(pte) & ~PAGE_MASK;
            printk("  PTE[%lu]: VA=0x%lx -> PA=0x%lx, prot=0x%lx\n", 
                   i, current_addr, pte_pa, pte_prot);
        } else {
            printk("  PTE[%lu]: VA=0x%lx -> UNMAPPED\n", i, current_addr);
        }
        i++;
        current_addr += PAGE_SIZE;
    } while (ptep++, current_addr != end);
}

// 打印PMD级页表内容  
static void print_pmd_level(pud_t *pudp, unsigned long addr, unsigned long end)
{
    pmd_t *pmdp = pmd_offset_phys(pudp, addr);
    unsigned long next;
    
    printk("PMD Level [0x%lx - 0x%lx]:\n", addr, end);
    do {
        next = pmd_addr_end(addr, end);
        pmd_t pmd = *pmdp;
        
        if (!pmd_none(pmd)) {
            if (pmd_sect(pmd)) {
                printk("  PMD Section: 0x%lx -> 0x%lx, prot: 0x%lx\n",
                       addr, pmd_val(pmd) & PMD_MASK, pmd_val(pmd) & ~PMD_MASK);
            } else {
                printk("  PMD Table: 0x%lx -> PTE base: 0x%lx\n",
                       addr, pmd_val(pmd) & PMD_MASK);
                // 递归打印PTE
                print_pte_level_simple(pmdp, addr, next);
            }
        }
    } while (pmdp++, addr = next, addr != end);
}

// 打印PUD级页表内容
static void print_pud_level(pgd_t *pgdp, unsigned long addr, unsigned long end)
{
    pud_t *pudp = pud_offset_phys(pgdp, addr);
    unsigned long next;
    
    printk("PUD Level [0x%lx - 0x%lx]:\n", addr, end);
    do {
        next = pud_addr_end(addr, end);
        pud_t pud = *pudp;
        
        if (!pud_none(pud)) {
            if (pud_sect(pud)) {
                printk("  PUD Section: 0x%lx -> 0x%lx\n", addr, pud_val(pud) & PUD_MASK);
            } else {
                printk("  PUD Table: 0x%lx -> PMD base: 0x%lx\n", 
                       addr, pud_val(pud) & PUD_MASK);
                // 递归打印PMD
                print_pmd_level(pudp, addr, next);
            }
        }
    } while (pudp++, addr = next, addr != end);
}

// 打印PGD级页表内容
void print_pgd_mapping(pgd_t *pgdir, unsigned long virt, unsigned long size)
{
    pgd_t *pgdp = pgd_offset_raw(pgdir, virt);
    unsigned long addr, end, next;
    
    addr = virt & PAGE_MASK;
    end = PAGE_ALIGN(virt + size);
    
    printk("=== Page Table Dump for [0x%lx - 0x%lx] ===\n", virt, end);
    
    do {
        next = pgd_addr_end(addr, end);
        pgd_t pgd = *pgdp;
        
        if (!pgd_none(pgd)) {
            if (pgd_sect(pgd)) {
                printk("PGD Section: 0x%lx -> 0x%lx\n", addr, pgd_val(pgd) & PGDIR_MASK);
            } else {
                printk("PGD Table: 0x%lx -> PUD base: 0x%lx\n", 
                       addr, pgd_val(pgd) & PGDIR_MASK);
                // 递归打印PUD
                print_pud_level(pgdp, addr, next);
            }
        }
    } while (pgdp++, addr = next, addr != end);
    
    printk("=== End of Page Table Dump ===\n");
}
/* debug end*/

static void alloc_init_pte(pmd_t *pmdp, unsigned long addr,
        unsigned long end, unsigned long phys,
        unsigned long prot,
        unsigned long (*alloc_pgtable)(void),
        unsigned long flags)
{
    pmd_t pmd = *pmdp;
    pte_t *ptep;
    unsigned long start_va = addr;
    unsigned long start_pa = phys;
    unsigned long count = 0;

    if (pmd_none(pmd)) {
        unsigned long pte_phys;
        pte_phys = alloc_pgtable();
        set_pmd(pmdp, __pmd(pte_phys | PMD_TYPE_TABLE));
        pmd = *pmdp;
        printk("Allocated PTE table: phys=0x%lx\n", pte_phys);
    }

    ptep = pte_offset_phys(pmdp, addr);
    do {
        set_pte(ptep, pfn_pte(phys >> PAGE_SHIFT, prot));
        count++;
        phys += PAGE_SIZE;
    } while (ptep++, addr += PAGE_SIZE, addr != end);

/*
    // 只打印PTE表的头尾信息
    if (count > 0) {
        unsigned long end_va = addr - PAGE_SIZE;  // 最后一个VA
        unsigned long end_pa = phys - PAGE_SIZE;  // 最后一个PA
        
        if (count == 1) {
            //printk("PTE: VA=0x%lx -> PA=0x%lx, prot=0x%lx\n", 
                   start_va, start_pa, prot);
        } else {
            printk("PTE: VA=0x%lx-0x%lx -> PA=0x%lx-0x%lx, prot=0x%lx [%lu pages]\n", 
                   start_va, end_va, start_pa, end_pa, prot, count);
        }
    }
	*/
}

void pmd_set_section(pmd_t *pmdp, unsigned long phys,
		unsigned long prot)
{
	unsigned long sect_prot = PMD_TYPE_SECT | mk_sect_prot(prot);

	pmd_t new_pmd = pfn_pmd(phys >> PMD_SHIFT, sect_prot);

	set_pmd(pmdp, new_pmd);
}

static void alloc_init_pmd(pud_t *pudp, unsigned long addr,
		unsigned long end, unsigned long phys,
		unsigned long prot,
		unsigned long (*alloc_pgtable)(void),
		unsigned long flags)
{
	pud_t pud = *pudp;
	pmd_t *pmdp;
	unsigned long next;

	if (pud_none(pud)) {
		unsigned long pmd_phys;

		pmd_phys = alloc_pgtable();
		set_pud(pudp, __pud(pmd_phys | PUD_TYPE_TABLE));
		pud = *pudp;
        printk("Allocated PMD table: phys=0x%lx\n", pmd_phys);

	}

	pmdp = pmd_offset_phys(pudp, addr);
	do {
		next = pmd_addr_end(addr, end);

		if (((addr | next | phys) & ~SECTION_MASK) == 0 &&
				(flags & NO_BLOCK_MAPPINGS) == 0){
			pmd_set_section(pmdp, phys, prot);
            //printk("PMD section mapping: va=0x%lx -> pa=0x%lx\n", addr, phys);
		}
		else
			alloc_init_pte(pmdp, addr, next, phys,
					prot,  alloc_pgtable, flags);

		phys += next - addr;
	} while (pmdp++, addr = next, addr != end);
}

static void alloc_init_pud(pgd_t *pgdp, unsigned long addr,
		unsigned long end, unsigned long phys,
		unsigned long prot,
		unsigned long (*alloc_pgtable)(void),
		unsigned long flags)
{
	pgd_t pgd = *pgdp;
	pud_t *pudp;
	unsigned long next;

	if (pgd_none(pgd)) {
		unsigned long pud_phys;

		pud_phys = alloc_pgtable();

		set_pgd(pgdp, __pgd(pud_phys | PUD_TYPE_TABLE));
		pgd = *pgdp;
	}

	pudp = pud_offset_phys(pgdp, addr);
	do {
		next = pud_addr_end(addr, end);
		alloc_init_pmd(pudp, addr, next, phys,
				prot, alloc_pgtable, flags);
		phys += next - addr;

	} while (pudp++, addr = next, addr != end);
}


static void __create_pgd_mapping(pgd_t *pgdir, unsigned long phys,
        unsigned long virt, unsigned long size,
        unsigned long prot,
        unsigned long (*alloc_pgtable)(void),
        unsigned long flags)
{
    pgd_t *pgdp = pgd_offset_raw(pgdir, virt);
    unsigned long addr, end, next;

    phys &= PAGE_MASK;
    addr = virt & PAGE_MASK;
    end = PAGE_ALIGN(virt + size);

    do{
        next = pgd_addr_end(addr, end);
        alloc_init_pud(pgdp, addr, next, phys,
                prot, alloc_pgtable, flags);
        phys += next - addr;
    } while (pgdp++, addr = next, addr != end);
}

unsigned long early_pgtable_alloc(void)
{
	unsigned long phys;

	phys = get_free_page();
	memset((void *)phys, 0, PAGE_SIZE);

	return phys;
}

static void create_identical_mapping(){
	unsigned long start;
	unsigned long end;

	start = (unsigned long)_text_boot;
	end = (unsigned long)_etext;
	__create_pgd_mapping((pgd_t *)idmap_pg_dir, start, start,
				end - start, PAGE_KERNEL_ROX,
				early_pgtable_alloc,
				0);

	/*map memory*/
	start = PAGE_ALIGN((unsigned long)_etext);
	end = TOTAL_MEMORY;
	__create_pgd_mapping((pgd_t *)idmap_pg_dir, start, start,
				end - start, PAGE_KERNEL,
				early_pgtable_alloc,
				0);
}

static void create_mmio_mapping(){
	__create_pgd_mapping((pgd_t *)idmap_pg_dir, PBASE, PBASE,
			DEVICE_SIZE, PROT_DEVICE_nGnRnE,
			early_pgtable_alloc,
			0);
}

static void cpu_init(void){
	unsigned long mair = 0;
	unsigned long tcr = 0;
	unsigned long tmp;
	unsigned long parang;

	asm("tlbi vmalle1");
	dsb(nsh);

	write_sysreg(3UL << 20, cpacr_el1);
	write_sysreg(1 << 12, mdscr_el1);

	mair = MAIR(0x00UL, MT_DEVICE_nGnRnE) |
	       MAIR(0x04UL, MT_DEVICE_nGnRE) |
	       MAIR(0x0cUL, MT_DEVICE_GRE) |
	       MAIR(0x44UL, MT_NORMAL_NC) |
	       MAIR(0xffUL, MT_NORMAL) |
	       MAIR(0xbbUL, MT_NORMAL_WT);
	write_sysreg(mair, mair_el1);
	tcr = TCR_TxSZ(VA_BITS) | TCR_TG_FLAGS;
	tmp = read_sysreg(ID_AA64MMFR0_EL1);
	parang = tmp & 0xf;
	if (parang > ID_AA64MMFR0_PARANGE_48)
		parang = ID_AA64MMFR0_PARANGE_48;

	tcr |= parang << TCR_IPS_SHIFT;

	write_sysreg(tcr, tcr_el1);

}

static int enable_mmu(void)
{
	unsigned long tmp;
	int tgran4;

	printk("Page table content at 0x%lx:\n", idmap_pg_dir);
	for (int i = 0; i < 10; i++) {
		printk("  [%d]: 0x%lx\n", i, *((uint64_t*)idmap_pg_dir + i));
	}

    // 打印开始信息
    printk("=== MMU Enable Start ===\n");
    
    // 打印当前MMU相关寄存器状态
    printk("Before configuration:\n");
    printk("  SCTLR_EL1 = 0x%lx\n", read_sysreg(sctlr_el1));
    printk("  TTBR0_EL1 = 0x%lx\n", read_sysreg(ttbr0_el1));
    printk("  TCR_EL1   = 0x%lx\n", read_sysreg(tcr_el1));
    printk("  MAIR_EL1  = 0x%lx\n", read_sysreg(mair_el1));
	printk("Current EL: 0x%lx\n", read_sysreg(CurrentEL));

	tmp = read_sysreg(ID_AA64MMFR0_EL1);
	tgran4 = (tmp >> ID_AA64MMFR0_TGRAN4_SHIFT) & 0xf;
    printk("ID_AA64MMFR0_EL1 = 0x%lx, tgran4 = %d\n", tmp, tgran4);
    
	if (tgran4 != ID_AA64MMFR0_TGRAN4_SUPPORTED)
		return -1;
	
    printk("Setting TTBR0_EL1 to 0x%p\n", idmap_pg_dir);
	write_sysreg(idmap_pg_dir, ttbr0_el1);
	isb();

    // 验证TTBR0设置
    printk("After TTBR0 set: 0x%lx\n", read_sysreg(ttbr0_el1));

    printk("Enabling MMU...\n");
	write_sysreg(SCTLR_ELx_M, sctlr_el1);
	isb();

    // 验证MMU是否启用
    //tmp = read_sysreg(sctlr_el1);
    //printk("After SCTLR set: 0x%lx (MMU %s)\n", tmp, (tmp & SCTLR_ELx_M) ? "ENABLED" : "DISABLED");

	asm("ic iallu");
	dsb(nsh);
	isb();

    printk("=== MMU Enable Complete ===\n");
	return 0;
}

void paging_init(void){
	printk("Setting TTBR0_EL1 to 0x%p-----------1\n", idmap_pg_dir);

	//memset(idmap_pg_dir, 0, PAGE_SIZE);

	create_identical_mapping();
	printk("test------3 \n");

	create_mmio_mapping();
	printk("test------4 \n");


    
    printk("=== Dumping MMIO mappings ===\n");
    //print_pgd_mapping((pgd_t *)idmap_pg_dir, PBASE, DEVICE_SIZE);

	//cpu_init();
	//printk("test------5 \n");

	//enable_mmu();
	printk("test------6 \n");


	printk("enable mmu done\n");
}